import winreg
import re
import datetime
import os
import csv

def filetime_to_datetime(ft):
    """Converts a Windows FILETIME to a Python datetime object."""
    us = (ft - 116444736000000000) // 10
    return datetime.datetime(1601, 1, 1) + datetime.timedelta(microseconds=us)

def get_usb_info_from_registry():
    """
    Extracts USB device information from the Windows Registry.
    Focuses on USBSTOR (mass storage) and general USB devices.
    """
    usb_devices = {} # Use a dictionary to store unique devices, key: (vid, pid, serial)

    # --- 1. USB Mass Storage Devices (USBSTOR) ---
    try:
        key_path_usbstor = r"SYSTEM\CurrentControlSet\Enum\USBSTOR"
        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, key_path_usbstor) as hkey_usbstor:
            for i in range(winreg.QueryInfoKey(hkey_usbstor)[0]):
                try:
                    subkey_name_disk = winreg.EnumKey(hkey_usbstor, i) # e.g., Disk&Ven_SanDisk&Prod_Cruzer_Blade&Rev_1.00
                    with winreg.OpenKey(hkey_usbstor, subkey_name_disk) as hkey_disk:
                        for j in range(winreg.QueryInfoKey(hkey_disk)[0]):
                            try:
                                subkey_name_instance = winreg.EnumKey(hkey_disk, j) # e.g., 20044352710192F06053&0
                                vid_match = re.search(r"Ven_([0-9A-F]{4})", subkey_name_disk, re.IGNORECASE)
                                pid_match = re.search(r"Prod_([0-9A-F]{4})", subkey_name_disk, re.IGNORECASE)
                                # Attempt to extract serial from instance ID (common pattern)
                                serial = subkey_name_instance.split('&')[0] if '&' in subkey_name_instance else subkey_name_instance
                                # For USBSTOR, the serial is often the part before the last '&' if present,
                                # or the whole string if no '&' (like some card readers' "slots")
                                # A better way for USBSTOR serials is often after the last '\' in the subkey_name_disk and before '&'
                                parts = subkey_name_disk.split('\\')
                                if len(parts) > 1:
                                    potential_serial_part = parts[-1]
                                    if '&' in potential_serial_part: # e.g., SomeProduct&Rev_1.00\SERIAL&0
                                        # this logic might need refinement based on actual registry structures
                                        pass # Serial for USBSTOR usually from instance key name
                                
                                vid = vid_match.group(1) if vid_match else "Unknown"
                                pid = pid_match.group(1) if pid_match else "Unknown"

                                device_key = (vid, pid, serial)
                                if device_key not in usb_devices:
                                    usb_devices[device_key] = {
                                        "vid": vid,
                                        "pid": pid,
                                        "serial_number": serial,
                                        "device_type": "USB Mass Storage",
                                        "description": "",
                                        "friendly_name": "",
                                        "manufacturer": "",
                                        "first_seen_setupapi": None,
                                        "last_seen_registry": None,
                                        "registry_path": f"{key_path_usbstor}\\{subkey_name_disk}\\{subkey_name_instance}"
                                    }
                                
                                with winreg.OpenKey(hkey_disk, subkey_name_instance) as hkey_instance:
                                    try:
                                        usb_devices[device_key]["description"] = winreg.QueryValueEx(hkey_instance, "DeviceDesc")[0]
                                    except FileNotFoundError:
                                        pass
                                    try:
                                        usb_devices[device_key]["friendly_name"] = winreg.QueryValueEx(hkey_instance, "FriendlyName")[0]
                                    except FileNotFoundError:
                                        pass
                                    try:
                                        usb_devices[device_key]["manufacturer"] = winreg.QueryValueEx(hkey_instance, "Mfg")[0]
                                    except FileNotFoundError:
                                        pass
                                    
                                    # Get LastWriteTime of the instance key as a proxy for last seen
                                    _, _, last_write_ft = winreg.QueryInfoKey(hkey_instance)
                                    usb_devices[device_key]["last_seen_registry"] = filetime_to_datetime(last_write_ft)

                            except WindowsError as e:
                                print(f"Error reading USBSTOR instance subkey: {e}")
                                continue
                except WindowsError as e:
                    print(f"Error reading USBSTOR disk subkey: {e}")
                    continue
    except FileNotFoundError:
        print(f"Registry key not found: {key_path_usbstor}")
    except WindowsError as e:
        print(f"Error accessing registry key {key_path_usbstor}: {e}")

    # --- 2. General USB Devices (USB) ---
    try:
        key_path_usb = r"SYSTEM\CurrentControlSet\Enum\USB"
        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, key_path_usb) as hkey_usb_root:
            for i in range(winreg.QueryInfoKey(hkey_usb_root)[0]):
                try:
                    subkey_name_vidpid = winreg.EnumKey(hkey_usb_root, i) # e.g., VID_04F2&PID_B107
                    vid_match = re.search(r"VID_([0-9A-F]{4})", subkey_name_vidpid, re.IGNORECASE)
                    pid_match = re.search(r"PID_([0-9A-F]{4})", subkey_name_vidpid, re.IGNORECASE)
                    vid = vid_match.group(1) if vid_match else "Unknown"
                    pid = pid_match.group(1) if pid_match else "Unknown"

                    with winreg.OpenKey(hkey_usb_root, subkey_name_vidpid) as hkey_vidpid:
                        for j in range(winreg.QueryInfoKey(hkey_vidpid)[0]):
                            try:
                                subkey_name_instance = winreg.EnumKey(hkey_vidpid, j) # This is often the serial number
                                serial = subkey_name_instance # For general USB, this is often the serial

                                device_key = (vid, pid, serial)
                                if device_key not in usb_devices:
                                    usb_devices[device_key] = {
                                        "vid": vid,
                                        "pid": pid,
                                        "serial_number": serial,
                                        "device_type": "General USB Device",
                                        "description": "",
                                        "friendly_name": "",
                                        "manufacturer": "",
                                        "first_seen_setupapi": None,
                                        "last_seen_registry": None,
                                        "registry_path": f"{key_path_usb}\\{subkey_name_vidpid}\\{subkey_name_instance}"
                                    }
                                
                                with winreg.OpenKey(hkey_vidpid, subkey_name_instance) as hkey_instance:
                                    try:
                                        usb_devices[device_key]["description"] = winreg.QueryValueEx(hkey_instance, "DeviceDesc")[0]
                                    except FileNotFoundError:
                                        pass
                                    try:
                                        usb_devices[device_key]["friendly_name"] = winreg.QueryValueEx(hkey_instance, "FriendlyName")[0]
                                    except FileNotFoundError:
                                        pass
                                    try:
                                        usb_devices[device_key]["manufacturer"] = winreg.QueryValueEx(hkey_instance, "Mfg")[0]
                                    except FileNotFoundError:
                                        pass
                                    
                                    _, _, last_write_ft = winreg.QueryInfoKey(hkey_instance)
                                    last_seen_dt = filetime_to_datetime(last_write_ft)
                                    
                                    # Update last_seen_registry if this one is newer or if it's a new entry
                                    if usb_devices[device_key]["last_seen_registry"] is None or \
                                       last_seen_dt > usb_devices[device_key]["last_seen_registry"]:
                                        usb_devices[device_key]["last_seen_registry"] = last_seen_dt
                                    
                                    # If it was previously marked as USBSTOR, keep it, otherwise update if not set
                                    if not usb_devices[device_key].get("device_type"):
                                        usb_devices[device_key]["device_type"] = "General USB Device"


                            except WindowsError as e:
                                print(f"Error reading USB instance subkey: {e}")
                                continue
                except WindowsError as e:
                    print(f"Error reading USB VID_PID subkey: {e}")
                    continue
    except FileNotFoundError:
        print(f"Registry key not found: {key_path_usb}")
    except WindowsError as e:
        print(f"Error accessing registry key {key_path_usb}: {e}")
        
    return usb_devices


def get_first_install_times_from_setupapi(devices_info):
    """
    Parses setupapi.dev.log for USB device installation timestamps.
    Updates the 'first_seen_setupapi' field in the devices_info dictionary.
    """
    log_file_path = os.path.join(os.environ.get("SystemRoot", "C:\\Windows"), "inf", "setupapi.dev.log")
    
    # Regex to find lines indicating device installation and capture VID, PID, Serial, and timestamp
    # Example line: >>>  [Device Install (Hardware IDs): USB\VID_0781&PID_5581\AA010304091047200170 - ROOT\USB\0000]
    # Example time line: >>>  Section start 2023/10/26 12:34:56.789
    # This regex is a bit simplified and might need adjustment
    device_install_regex = re.compile(r"USB\\VID_([0-9A-F]{4})&PID_([0-9A-F]{4})(?:\\|&)([^\s\]]+)")
    timestamp_regex = re.compile(r">>>\s*Section start\s*(\d{4}/\d{2}/\d{2}\s*\d{2}:\d{2}:\d{2}\.\d{3})")
    # More specific for device install start:
    # >>>  [Device Install (Hardware IDs): USB\VID_0781&PID_5583\20060877600D1B116E5A]
    # ...
    # <<<  Section end 2023/03/03 10:43:57.208
    # <<<  [Exit status: SUCCESS]
    
    current_timestamp = None

    try:
        with open(log_file_path, 'r', encoding='utf-8', errors='ignore') as f:
            for line in f:
                ts_match = timestamp_regex.search(line)
                if ts_match:
                    try:
                        current_timestamp = datetime.datetime.strptime(ts_match.group(1), "%Y/%m/%d %H:%M:%S.%f")
                    except ValueError:
                        # Handle cases where millisecond might be missing or format slightly different
                        try:
                            current_timestamp = datetime.datetime.strptime(ts_match.group(1).split('.')[0], "%Y/%m/%d %H:%M:%S")
                        except ValueError:
                            current_timestamp = None # In case of parsing error

                if "Device Install (Hardware IDs)" in line and current_timestamp:
                    match = device_install_regex.search(line)
                    if match:
                        vid = match.group(1).upper()
                        pid = match.group(2).upper()
                        # Serial might have extra info after '&', remove it
                        serial_full = match.group(3).upper()
                        serial = serial_full.split('&')[0] 

                        device_key = (vid, pid, serial)
                        
                        if device_key in devices_info:
                            # Only update if this timestamp is earlier (first install)
                            if devices_info[device_key]["first_seen_setupapi"] is None or \
                               current_timestamp < devices_info[device_key]["first_seen_setupapi"]:
                                devices_info[device_key]["first_seen_setupapi"] = current_timestamp
                        else:
                            # Device found in setupapi but not (yet) in registry scan. Add it.
                            # This can happen for devices that were connected but perhaps drivers not fully enumerated
                            # or if registry parsing missed it.
                            devices_info[device_key] = {
                                "vid": vid,
                                "pid": pid,
                                "serial_number": serial,
                                "device_type": "From SetupAPI Log",
                                "description": "N/A (from log)",
                                "friendly_name": "N/A (from log)",
                                "manufacturer": "N/A (from log)",
                                "first_seen_setupapi": current_timestamp,
                                "last_seen_registry": None,
                                "registry_path": "N/A (from log)"
                            }

    except FileNotFoundError:
        print(f"setupapi.dev.log not found at {log_file_path}")
    except Exception as e:
        print(f"Error reading or parsing setupapi.dev.log: {e}")


def main():
    print("Attempting to gather USB device information. This requires Administrator privileges.")
    
    # 1. Get info from Registry
    usb_devices_data = get_usb_info_from_registry()
    
    # 2. Get first install times from setupapi.dev.log and merge
    get_first_install_times_from_setupapi(usb_devices_data)

    print(f"\nFound {len(usb_devices_data)} unique USB device entries.\n")

    output_data = []
    for key, info in usb_devices_data.items():
        output_data.append({
            "VID": info["vid"],
            "PID": info["pid"],
            "Serial Number": info["serial_number"],
            "Device Type": info["device_type"],
            "Description": info["description"],
            "Friendly Name": info["friendly_name"],
            "Manufacturer": info.get("manufacturer", "N/A"), # Use .get for safety
            "First Seen (SetupAPI)": info["first_seen_setupapi"].strftime('%Y-%m-%d %H:%M:%S') if info["first_seen_setupapi"] else "N/A",
            "Last Known (Registry)": info["last_seen_registry"].strftime('%Y-%m-%d %H:%M:%S') if info["last_seen_registry"] else "N/A",
            "Registry Path": info.get("registry_path", "N/A")
        })
        
        print(f"VID: {info['vid']}, PID: {info['pid']}, Serial: {info['serial_number']}")
        print(f"  Type: {info['device_type']}")
        if info['description']: print(f"  Description: {info['description']}")
        if info['friendly_name']: print(f"  Friendly Name: {info['friendly_name']}")
        if info.get('manufacturer'): print(f"  Manufacturer: {info.get('manufacturer')}")
        print(f"  First Seen (SetupAPI): {info['first_seen_setupapi'].strftime('%Y-%m-%d %H:%M:%S') if info['first_seen_setupapi'] else 'N/A'}")
        print(f"  Last Known (Registry): {info['last_seen_registry'].strftime('%Y-%m-%d %H:%M:%S') if info['last_seen_registry'] else 'N/A'}")
        if info.get('registry_path'): print(f"  Registry Path: {info.get('registry_path')}")
        print("-" * 30)

    # Optional: Save to CSV
    if output_data:
        csv_file = "usb_device_history.csv"
        print(f"\nSaving data to {csv_file}...")
        try:
            with open(csv_file, 'w', newline='', encoding='utf-8') as f:
                writer = csv.DictWriter(f, fieldnames=output_data[0].keys())
                writer.writeheader()
                writer.writerows(output_data)
            print(f"Successfully saved to {csv_file}")
        except IOError as e:
            print(f"Error saving CSV: {e}")

if __name__ == "__main__":
    # Check for admin rights (simple check, might not be foolproof)
    try:
        is_admin = (os.getuid() == 0) # For Linux/macOS
    except AttributeError:
        # For Windows
        import ctypes
        try:
            is_admin = ctypes.windll.shell32.IsUserAnAdmin() != 0
        except Exception:
            is_admin = False # Assume not admin if check fails

    if not is_admin:
        print("Warning: This script likely requires administrator privileges to access all registry keys and system logs.")
        print("Please try running it as an administrator if you encounter errors or missing data.")
        # For Windows, you could attempt to re-launch as admin here, but it's complex.
        # For a course project, simply instructing the user is usually sufficient.

    main()